import sys, os
import torch
import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg
import matplotlib.pyplot as plt
from PIL import Image
import imageio

current_dir = os.path.dirname(__file__)
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.append(parent_dir)
from config_env import get_env_config
from config_true_env import get_config
from train.separated_model import Separated_model
from envs.env_wrappers import DummyVecEnv
parser_ev = get_env_config()

def parse_args(args, parser):
    parser.add_argument("--scenario_name", type=str, default="MyEnv_true", help="Which scenario to run on")
    parser.add_argument("--num_landmarks", type=int, default=3)
    parser.add_argument("--num_agents", type=int, default=parser_ev["agent_num"], help="number of players")
    all_args = parser.parse_known_args(args)[0]
    return all_args

def make_eval_env(env_size = []):
    def get_env_fn(rank):
        def init_env(env_size):
            from envs.env_continuous import ContinuousActionEnv
            env = ContinuousActionEnv(env_size)
            env.seed(1 + rank * 1000)
            return env
        return init_env
    return DummyVecEnv([get_env_fn(i) for i in range(1)], env_size)

def visualize_environment(obs, num_agents, save_path=''):
    plt.clf()  # 清除当前图像
    # 提取 x 和 y 坐标
    agent_point = []
    for ag in range(num_agents):
        if parser_ev["obs_dim"] == 16:
            agent_point.append([coord[ag][4*ag:4*ag+2] for coord in obs])
        elif parser_ev["obs_dim"] == 8:
            agent_point.append([coord[ag][0:2] for coord in obs])

        # 绘制散点图
        plt.scatter([item[0] for item in agent_point[ag]], [item[1] for item in agent_point[ag]], marker='o', label=f'Agent {ag + 1}')
        for point in agent_point[ag]:
            circle = plt.Circle((point[0], point[1]), radius=parser_ev["Radius_collision"], fill=False, edgecolor='blue', linestyle='--')
            plt.gca().add_patch(circle)
    # print(obs)
    # print("agent_point",agent_point)
    # print("agent_point[ag][:][0]",agent_point[ag])
    # # 绘制圆
    plt.scatter(obs[0][0][-4], obs[0][0][-3], marker='o', label=f'Target')
    circle = plt.Circle((obs[0][0][-4], obs[0][0][-3]), radius=parser_ev["Radius_detect"]*0.75, color='red', fill=False)
    circle2 = plt.Circle((obs[0][0][-4], obs[0][0][-3]), radius=parser_ev["Radius_detect"]*1.25, color='red', fill=False)
    plt.gca().add_patch(circle)
    plt.gca().add_patch(circle2)
    # 添加图例
    plt.legend()
    # plt.axis('equal')  # 保持坐标轴比例一致
    # 添加标签和标题
    plt.xlim(-parser_ev["max_map"],parser_ev["max_map"])
    plt.ylim(-parser_ev["max_map"],parser_ev["max_map"])
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.title('Environment Visualization')
    # 将plt转化为numpy数据
    canvas = FigureCanvasAgg(plt.gcf())
    # # 绘制图像
    canvas.draw()
    # 获取图像尺寸
    w, h = canvas.get_width_height()
    # 解码string 得到argb图像
    buf = np.fromstring(canvas.tostring_argb(), dtype=np.uint8)
 
    # 重构成w h 4(argb)图像
    buf.shape = (w, h, 4)
    # 转换为 RGBA
    buf = np.roll(buf, 3, axis=2)
    # 得到 Image RGBA图像对象 (需要Image对象的同学到此为止就可以了)
    image = Image.frombytes("RGBA", (w, h), buf.tostring())
    # 转换为numpy array rgba四通道数组
    image = np.asarray(image)
    # 转换为rgb图像
    rgb_image = image[:, :, :3]
    
    # # 保存图像到文件
    # plt.savefig(save_path)
    # plt.close()
    # plt.pause(0.00001)  # 暂停一会以更新图像

    return rgb_image

if __name__ == "__main__":
    parser = get_config()
    all_args = parse_args(sys.argv[1:], parser)

    # cuda
    if all_args.cuda and torch.cuda.is_available():
        print("choose to use gpu...")
        device = torch.device("cuda")
        torch.set_num_threads(all_args.n_training_threads)
        if all_args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        print("choose to use cpu...")
        device = torch.device("cpu")
        torch.set_num_threads(all_args.n_training_threads)

    all_args.model_dir = parser_ev["eval_model_dir"]
    config = {
        "all_args": all_args,
        "num_agents":all_args.num_agents,
        "device": device,
    }
    agents = []
    for i in range(3):
        agent = Separated_model(config, 0)
        agents.append(agent)
    env = make_eval_env(["Circle", [0, 0], 10])
    obs = env.reset()
    images = []
    
    for i in range(200):
        actions = []
        for i in range(3):
            action = agents[i].step(obs[:,i])[0][0]
            actions.append(action)
        # actions增加一维
        actions = np.expand_dims(actions, axis=0)
        image = visualize_environment(obs, 3)
        images.append(image)
        obs, eval_communicate_mask, eval_rewards, eval_dones, eval_infos  = env.step(actions)
    imageio.mimsave( 'D:/Reinforce/light_mappo-main/light_mappo-main/True_Model_Env/episode.gif', \
                    images, 'GIF', duration=1,  overwrite=True)